# -*- coding: utf-8 -*-
import numpy as np

out_file = "../skynet_flow.h"
bitw, sym_bitw = 5,  "BIT_WEIT"
bitb, sym_bitb = 16, "BIT_BIAS"
bitm, sym_bitm = 16, "BIT_MULT"

def int2hex(x):
	return '\"{:#x}\"'.format(x)

def write_const(file, symbol, val):
	code = "constexpr unsigned {} = {:d};\n".format(symbol, val)
	file.write(code)
	return

def write_line(file, data, bit):
	assert len(data.shape) == 2
	data[data < 0] += np.power(2, bit)
	outer, inner = data.shape
	line = [None] * outer
	file.write("\t")
	for i in range(outer):
		cat = 0
		for j in range(inner):
			cat |= int(data[i, j]) << (bit * j)
		line[i] = cat
	file.write(", ".join(map(int2hex, line)))
	file.write(",\n")
	return

def write_arr1D(file, data, bit, sym_arr, sym_width, sym_dim):
	sp = data.shape
	assert len(sp) == 2
	if sp[0] % 8 == 0:
		data = data.reshape(sp[0] // 8, 8, sp[1])
	else:
		data = data.reshape(1, sp[0], sp[1])
	title = "const ap_uint<{0}> {1}[{2}] = {{\n".format(sym_width, sym_arr, sym_dim)
	file.write(title)
	for i in range(data.shape[0]):
		write_line(file, data[i], bit)
	file.write("};\n")
	return

def write_arr2D(file, data, bit, sym_arr, sym_width, sym_dim1, sym_dim2):
	assert len(data.shape) == 3
	title = "const ap_uint<{0}> {1}[{2}][{3}] = {{\n".format(sym_width, sym_arr, sym_dim1, sym_dim2)
	file.write(title)
	for i in range(data.shape[0]):
		write_line(file, data[i], bit)
	file.write("};\n")
	return

def write_DW_weight(file, weight, bias, mult, n_io, n_act, tag):
	print(weight.shape)
	n_ch = weight.shape[2]
	assert len(weight.shape) == 3
	assert weight.shape[0] == weight.shape[1] == 3
	assert len(bias.shape) == len(mult.shape) == 1
	assert bias.shape[0] == mult.shape[0] == n_ch
	assert n_ch >= n_io and n_ch % n_io == 0
	assert n_ch >= n_act and n_ch % n_act == 0
	
	data = weight.reshape(9, n_ch // n_io, n_io).transpose(1, 0, 2)
	assert data.shape == (n_ch // n_io, 9, n_io)
	# data[data < 0] += np.power(2, bitw)
	
	sym_arr = "{}_DW".format(tag)
	sym_nch = "{}_NCH".format(sym_arr)
	sym_nio = "{}_NIO".format(sym_arr)
	sym_act = "{}_NACT".format(sym_arr)
	write_const(file, sym_nch, n_ch)
	write_const(file, sym_nio, n_io)
	write_const(file, sym_act, n_act)
   
	sym_width = "{} * {}".format(sym_nio, sym_bitw)
	sym_dim = "{} / {}".format(sym_nch, sym_nio)
	write_arr2D(file, data, bitw, sym_arr, sym_width, sym_dim, 9)
	
	bias = bias.reshape(n_ch // n_act, n_act)
	sym_arr = "{}_DB".format(tag)
	sym_width = "{} * {}".format(sym_act, sym_bitb)
	sym_dim = "{} / {}".format(sym_nch, sym_act)
	write_arr1D(file, bias, bitb, sym_arr, sym_width, sym_dim)
	
	mult = mult.reshape(n_ch // n_act, n_act)
	sym_arr = "{}_DM".format(tag)
	sym_width = "{} * {}".format(sym_act, sym_bitm)
	sym_dim = "{} / {}".format(sym_nch, sym_act)
	write_arr1D(file, mult, bitm, sym_arr, sym_width, sym_dim)
	
	file.write("\n")
	return

def write_DWEK_weight(file, weight, bias, mult, n_io, n_act, tag):
	print(weight.shape)
	n_ch = weight.shape[2]
	assert len(weight.shape) == 3
	assert weight.shape[0] == weight.shape[1] == 3
	assert len(bias.shape) == len(mult.shape) == 1
	assert bias.shape[0] == mult.shape[0] == n_ch
	assert n_ch >= n_io and n_ch % n_io == 0
	assert n_ch >= n_act and n_ch % n_act == 0
	
	kh, kw = 3, 3

	wt_t = weight.reshape(kh, kw, n_ch // n_io, n_io)
	wt_t = wt_t.transpose(2, 0, 1, 3)
	assert wt_t.shape == (n_ch // n_io, kh, kw, n_io)
	wt_t = wt_t.reshape(n_ch // n_io, kh, kw*n_io)

	sym_arr = "{}_DW".format(tag)
	sym_nch = "{}_NCH".format(sym_arr)
	sym_nio = "{}_NIO".format(sym_arr)
	sym_act = "{}_NACT".format(sym_arr)
	sym_k   = "{}_K".format(sym_arr)
	write_const(file, sym_nch, n_ch)
	write_const(file, sym_nio, n_io)
	write_const(file, sym_act, n_act)
	write_const(file, sym_k, kw)

	sym_width = "{} * {} * {}".format(sym_nio, sym_k, sym_bitw)
	sym_dim1 = "{} / {}".format(sym_nch, sym_nio)
	sym_dim2 = "{}".format(sym_k)
	write_arr2D(file, wt_t, bitw, sym_arr, sym_width, sym_dim1, sym_dim2)

	bias = bias.reshape(n_ch // n_act, n_act)
	sym_arr = "{}_DB".format(tag)
	sym_width = "{} * {}".format(sym_act, sym_bitb)
	sym_dim = "{} / {}".format(sym_nch, sym_act)
	write_arr1D(file, bias, bitb, sym_arr, sym_width, sym_dim)
	
	mult = mult.reshape(n_ch // n_act, n_act)
	sym_arr = "{}_DM".format(tag)
	sym_width = "{} * {}".format(sym_act, sym_bitm)
	sym_dim = "{} / {}".format(sym_nch, sym_act)
	write_arr1D(file, mult, bitm, sym_arr, sym_width, sym_dim)
	
	file.write("\n")
	return

def write_PW_weight(file, weight, bias, mult, n_in, n_out, n_act, tag):
	n_och, n_ich = weight.shape
	assert len(weight.shape) == 2
	assert len(bias.shape) == len(mult.shape) == 1
	assert bias.shape[0] == mult.shape[0] == n_och
	assert n_och >= n_out and n_och % n_out == 0
	assert n_ich >= n_in  and n_ich % n_in  == 0
	if n_act > 0:
		assert n_och >= n_act and n_och % n_act == 0
	
	data = weight.reshape(n_och // n_out, n_out, n_ich // n_in, n_in).transpose(0, 2, 1, 3)
	data = data.reshape(n_och // n_out, n_ich // n_in, n_out * n_in)
	# data[data < 0] += np.power(2, bitw)
	
	sym_arr = "{}_PW".format(tag)
	sym_ich = "{}_NICH".format(sym_arr)
	sym_och = "{}_NOCH".format(sym_arr)
	sym_in  =  "{}_NIN".format(sym_arr)
	sym_out = "{}_NOUT".format(sym_arr)
	sym_act = "{}_NACT".format(sym_arr)
	write_const(file, sym_ich, n_ich)
	write_const(file, sym_och, n_och)
	write_const(file, sym_in , n_in )
	write_const(file, sym_out, n_out)
	if n_act > 0:
		write_const(file, sym_act, n_act)
	
	sym_width = "{} * {} * {}".format(sym_out, sym_in, sym_bitw)
	sym_dim = "{} / {}".format(sym_och, sym_out)
	sym_dim2 = "{} / {}".format(sym_ich, sym_in)
	write_arr2D(file, data, bitw, sym_arr, sym_width, sym_dim, sym_dim2)
	
	if n_act > 0:
		bias = bias.reshape(n_och // n_act, n_act)
		sym_arr = "{}_PB".format(tag)
		sym_width = "{} * {}".format(sym_act, sym_bitb)
		sym_dim = "{} / {}".format(sym_och, sym_act)
		write_arr1D(file, bias, bitb, sym_arr, sym_width, sym_dim)
		
		mult = mult.reshape(n_och // n_act, n_act)
		sym_arr = "{}_PM".format(tag)
		sym_width = "{} * {}".format(sym_act, sym_bitm)
		sym_dim = "{} / {}".format(sym_och, sym_act)
		write_arr1D(file, mult, bitm, sym_arr, sym_width, sym_dim)
	
	file.write("\n")
	return

def write_BM_weight(file, bias, mult, n_io, symbol):
	return

weight = np.load("skynet_weight_int32_KKC_OCIC.npz")

with open(out_file, "w") as f:
	f.write("#pragma once\n")
	f.write("#include<fstream>\n")
	f.write("#include<iostream>\n")
	f.write("#include<string>\n")
	f.write("#include<cmath>\n")
	f.write("#include \"stream_tools.h\"\n\n")
	
	f.write("constexpr unsigned N_BATCH = 1;\n")
	f.write("constexpr unsigned N_IN = 4;\n")
	f.write("constexpr unsigned N_OUT = 1;\n")
	f.write("constexpr unsigned BIT_ACTV = 8;\n")
	f.write("constexpr unsigned BIT_WEIT = 5;\n")
	f.write("constexpr unsigned BIT_BIAS = 16;\n")
	f.write("constexpr unsigned BIT_MULT = 16;\n")
	f.write("constexpr unsigned BIT_CONV = 16;\n")
	f.write("constexpr unsigned BIT_NORM = 18;\n")
	f.write("constexpr unsigned R_SHIFT = 16;\n")
	f.write("constexpr unsigned ROWIN = 360;\n")
	f.write("constexpr unsigned COLIN = 640;\n")
	f.write("constexpr unsigned ROW0 = 160;\n")
	f.write("constexpr unsigned COL0 = 320;\n")
	f.write("constexpr unsigned ROW1 = ROW0 / 2;\n")
	f.write("constexpr unsigned COL1 = COL0 / 2;\n")
	f.write("constexpr unsigned ROW2 = ROW1 / 2;\n")
	f.write("constexpr unsigned COL2 = COL1 / 2;\n")
	f.write("constexpr unsigned ROW3 = ROW2 / 2;\n")
	f.write("constexpr unsigned COL3 = COL2 / 2;\n")
	f.write("constexpr unsigned DEFAULT_DEPTH = 4;\n\n")
	
	# Bundle 1
	write_DWEK_weight(f, weight["L0-W"], weight["L0-B"], weight["L0-M"], 3, 1, "L0")
	write_PW_weight(f, weight["L1-W"], weight["L1-B"], weight["L1-M"], 3, 12, 12, "L0")
	
	# Bundle 2
	write_DWEK_weight(f, weight["L2-W"], weight["L2-B"], weight["L2-M"], 12, 3, "L1")
	write_PW_weight(f, weight["L3-W"], weight["L3-B"], weight["L3-M"], 12, 24, 8, "L1")
	
	# Bundle 3
	write_DW_weight(f, weight["L4-W"], weight["L4-B"], weight["L4-M"], 16, 2, "L2")
	write_PW_weight(f, weight["L5-W"], weight["L5-B"], weight["L5-M"], 24, 12, 4, "L2")
	
	# Bundle 4
	write_DW_weight(f, weight["L6-W"], weight["L6-B"], weight["L6-M"], 8, 1, "L3")
	write_PW_weight(f, weight["L7-W"], weight["L7-B"], weight["L7-M"], 24, 12, 2, "L3")
	
	# Bundle 5
	write_DW_weight(f, weight["L8-W"], weight["L8-B"], weight["L8-M"], 16, 2, "L4")
	write_PW_weight(f, weight["L9-W"], weight["L9-B"], weight["L9-M"], 24, 32, 2, "L4")

	# Bundle 6
	write_DWEK_weight(f, weight["L10-W"], weight["L10-B"], weight["L10-M"], 16, 8, "L5")
	write_PW_weight(f, weight["L11-W"], weight["L11-B"], weight["L11-M"], 20, 24, 1, "L5")
	
	# Bundle 7
	write_PW_weight(f, weight["L12-W"], weight["L12-B"], weight["L12-M"], 4, 1, -1, "L6")


	f.write("void skynet_flow(data_stream<N_IN * BIT_ACTV>& in, axiu_stream<N_OUT * BIT_CONV>& out);")
